import numpy as np
import random, time
from scipy.integrate import ode
from UTILS.tensor_ops import distance_matrix, repeat_at, delta_matrix
from VISUALIZE.mcom import mcom
PI = np.pi
def run():
    # 可视化界面初始化
    可视化桥 = mcom(path='RECYCLE/v2d_logger/', draw_mode='Threejs')
    可视化桥.初始化3D(); 可视化桥.设置样式('star')
    可视化桥.其他几何体之旋转缩放和平移('ball', 'SphereGeometry(1)',   0,0,0,  1,1,1, 0,0,0) # 球体   
    # 设置初始参数
    # colors = ['Yellow', 'Red', 'Blue', 'White','Yellow', 'Red', 'Blue', 'White','Yellow', 'Red', 'Blue', 'White','Yellow', 'Red', 'Blue', 'White','Yellow', 'Red', 'Blue', 'White']
    colors = ['Yellow', 'Red']
    N体 = 69
    theta = 2*PI/(N体-1)
    相位 = [i*theta for i in range(N体-1)] + [0]
    位置 = [[np.cos(相位[i]), np.sin(相位[i]), (-1)**i] for i in range(N体-1)] + [[0,0,0]]; 位置 = np.array(位置)*20
    速度 = [[np.cos(相位[i]+np.pi/2), np.sin(相位[i]+np.pi/2), 0] for i in range(N体-1)] + [[0,0,0]]; 速度 = np.array(速度)*7
    质量 = [1.5 for _ in range(N体-1)] + [200]; 质量 = np.array(质量)
    G = 16 # 常量 G
    N体=len(位置)
    assert N体==len(位置) and N体==len(速度)
    dT =  0.001 # 时间间隔

    def pack_y(位置, 速度):
        return np.concatenate((位置.reshape(-1), 速度.reshape(-1)))

    def unpack_y(y):
        位置, 速度  = y.reshape(2, N体, 3)
        return 位置, 速度

    def f(t, y):
        位置, 速度 = unpack_y(y)
        # F = (GMm)/距离矩阵^2
        距离矩阵 = distance_matrix(位置) # 计算距离矩阵
        距离矩阵平方 = 距离矩阵 * 距离矩阵.T
        # 计算Mm矩阵：
        质量矩阵 = repeat_at(tensor=质量, insert_dim=-1, n_times=N体)   
        Mm矩阵 = 质量矩阵 * 质量矩阵.T
        # 计算引力
        引力标量矩阵 = G*Mm矩阵/(距离矩阵平方+1e-10)
        # 对角线元素清零
        for i in range(N体): 引力标量矩阵[i,i] = 0
        # 引力的方向，即位移的方向
        引力方向 = delta_matrix(位置) # 形状为 (N体, N体, 3)
        # 引力方向归一化，乘以引力矩阵（标量）
        引力方向 = 引力方向/(np.linalg.norm(引力方向,axis=-1, keepdims=True)+1e-10)
        引力 = 引力方向*repeat_at(tensor=引力标量矩阵, insert_dim=-1, n_times=3)
        引力合力 = 引力.sum(1)
        # a=F/m
        加速度 = 引力合力/repeat_at(tensor=质量, insert_dim=-1, n_times=3)
        位置导数 = 速度
        速度导数 = 加速度
        y_dot = pack_y(位置导数, 速度导数)
        return y_dot

    r = ode(f).set_integrator('dop853') # lsoda is fast, dop853 is precise but slow!
    y0 = pack_y(位置, 速度);    r.set_initial_value(y0, 0)
    while r.successful() and r.t < 100:
        时间 = int((r.t+dT)/dT)
        位置, 速度 = unpack_y(r.integrate(r.t+dT))
        if 时间%10==0:
            位置_实部 = 位置.real
            for body in range(N体):
                可视化桥.发送几何体('ball|%d|%s|%.2f'%(body, colors[body%2], 0.4),  # 填入核心参量： “已声明的形状|几何体之ID标识|颜色|整体大小”
                    位置_实部[body, 0], 位置_实部[body, 1], 位置_实部[body, 2], ro_x=0, ro_y=0, ro_z=0, # 6 DOF
                    opacity=1, label='', label_color='white',  track_n_frame=2500,  track_tension=0, track_color=colors[body%2],    # 轨迹的颜色显示，输入js颜色名或者hex值均可
                    )
            可视化桥.结束关键帧()
        print('\r %.2f'%(r.t+dT), end='', flush=True)



DATA_ARR = [(0.3471168881,	0.5327249454),
(0.3068934205,	0.1255065670),
(0.6150407229,	0.5226158545),
(0.5379557207,	0.3414578545),
(0.4112926910,	0.2607551013),
(0.4425908552,	0.4235138348),
(0.1214534165,	0.1012023800),
(0.4094945913,	0.3628231655),
(0.5255769251,	0.2501253528),
(0.4121028725,	0.2833837497),
(0.4364192674,	0.4457095477),
(0.1182009612,	0.5847367662),
(0.4763527642,	0.3789434497),
(0.4027121690,	0.2100155085),
(0.3960494651,	0.3529413931),
(0.4340543928,	0.4606927255),
(0.1842784887,	0.5871881740),
(0.4580378828,	0.4093753024),
(0.4152505707,	0.2913461808),
(0.0969422653,	0.5615968396),
(0.4082108156,	0.2436851904),
(0.4289870678,	0.4723797648),
(0.4138807521,	0.3477955684),
(0.4905050535,	0.4044215155),
(0.1982770999,	0.5760625510),
(0.3991287659,	0.1847081193),
(0.4151691260,	0.2953409098),
(0.0490506729,	0.5901941115),
(0.3447503346,	0.3930446386),
(0.2203123981,	0.5718227262),
(0.4441311511,	0.3823106072),
(0.4041322114,	0.2191641945),
(0.4089912496,	0.3457125084),
(0.4572423635,	0.4074171858),
(0.3972193787,	0.1691985386),
(0.4160674674,	0.2971499303),
(0.1038901209,	0.5907210858),
(0.4135366646,	0.2710056059),
(0.3967429300,	0.3708809839),
(0.4073762182,	0.2388431491),
(0.1334658448,	0.5838277552),
(0.4015585101,	0.2022664462),
(0.4157753114,	0.2981866614),
(0.4188475683,	0.2770588072),
(0.3960577146,	0.1586009694),
(0.4096220179,	0.2516961122),
(0.4046803380,	0.2225972157),
(0.2744365569,	0.3448729405),
(0.5196830772,	0.2059463323),
(0.0481112024,	0.5584284134),
(0.3240593189,	0.3665712251),
(0.4392326535,	0.4490130778),
(0.4157392818,	0.3443683086),
(0.2591633491,	0.5655538324),
(0.4372549674,	0.4568896400),
(0.4713347232,	0.3856148608),
(0.1538860573,	0.5805527776),
(0.4162411626,	0.4666456940),
(0.4602603277,	0.4047836526),
(0.4260682804,	0.4689275280),
(0.4115604103,	0.3439515057),
(0.4538585041,	0.4151335211),
(0.1807391308,	0.5804273075),
(0.0616295799,	0.5887684684),
(0.4164511851,	0.2985460701),
(0.4755802612,	0.3834289883),
(0.4189618506,	0.3575723235),
(0.3998065241,	0.1898070007),
(0.1938115742,	0.5776286093),
(0.4311833725,	0.4733759409),
(0.0960540551,	0.5889671026),
(0.3952890968,	0.1508523182),
(0.4395870869,	0.4502293684),
(0.3674149486,	0.3838228297),
(0.3939961214,	0.3865732823),
(0.3430511121,	0.3718832371),
(0.4159559963,	0.2988672319),
(0.4069867693,	0.2365548555),
(0.4589974678,	0.4050034288),
(0.4139140441,	0.2867091579),
(0.4379684294,	0.4556963530),
(0.4126145620,	0.3435200139),
(0.4561844518,	0.4117892184),
(0.3985475384,	0.1801838327),
(0.4126594494,	0.2674571676),
(0.3921967782,	0.3732670713),
(0.4180525186,	0.2985930530),
(0.1382080826,	0.5834341742),
(0.3974576989,	0.2491169363),
(0.0188646524,	0.5900765061),
(0.3947500473,	0.1449169673),
(0.4146876683,	0.2893277049),
(0.4049707755,	0.2243966352),
(0.4948547469,	0.3816681253),
(0.4098282065,	0.4859736518),
(0.3557656112,	0.2101434426),
(0.1458080611,	0.5760627737),
(0.4373040972,	0.1890148898),
(0.3693848766,	0.2787954099),
(0.0670760777,	0.5889627892),
(0.4024084539,	0.3670928188),
(0.4101448519,	0.2545904027),
(0.3976053675,	0.1724965035),
(0.4070387234,	0.3507052842),
(0.3654211421,	0.1497581124),
(0.3288064734,	0.2480300294),
(0.0921822548,	0.5888237645),
(0.3943552252,	0.1402131694),
(0.4164334853,	0.2763211125),
(0.4033838433,	0.2143907433),
(0.1818464802,	0.5795638348),
(0.4109222799,	0.3490743143),
(0.4159829843,	0.2928634271),
(0.3968779003,	0.1661961334),
(0.3369172422,	0.2901238678),
(0.4786094582,	0.2027583452),
(0.1262212709,	0.5844212753),
(0.3596345346,	0.3768544947),
(0.3735693647,	0.3814544097),
(0.3940559464,	0.1363870337),
(0.0398375597,	0.5900216545),
(0.3366050958,	0.1545246545),
(0.4148912516,	0.3594943116),
(0.4021066869,	0.2059880758),
(0.2984912721,	0.2642719294),
(0.3151887676,	0.2026092096),
(0.3682819786,	0.0540125357),
(0.4094064944,	0.2504917575),
(0.3963021063,	0.1609272065),
(0.4858230168,	0.1897852046),
(0.3938225013,	0.1332100629),
(0.2844202818,	0.2576751665),
(0.3613650149,	0.0662137880),
(0.4052729644,	0.2262551267),
(0.2792503345,	0.2309398618),
(0.4088863339,	0.2475714173),
(0.2755037979,	0.2512907499),
(0.3000356874,	0.1954805934),
(0.2759217512,	0.2671023857),
(0.3060236507,	0.1574877923),
(0.2738106556,	0.2266665399),
(0.2716635736,	0.2803602873),
(0.2690738475,	0.2460748830),
(0.2813047676,	0.2931043691),
(0.2844033114,	0.2864772055),
(0.3406159593,	0.0878815592),
(0.2745076841,	0.2185259293),
(0.5236560165,	0.2615293843),
(0.5252433392,	0.2395323622),
(0.2485711669,	0.2285602706),
(0.4717436769,	0.4159399383),
(0.4244996083,	0.4906100936),
(0.3303010815,	0.4704859697),
(0.2538439618,	0.5599410925),
(0.4824594978,	0.3739325421),
(0.4299480127,	0.4772693278),
(0.2242029652,	0.5716643145),
(0.4721688158,	0.3825825184),
(0.2338958736,	0.5689557210),
(0.4067612870,	0.2352215102),
(0.4664129400,	0.2183545493),
(0.1082811863,	0.5839631505),
(0.3733531828,	0.0372160210),
(0.3289795673,	0.2068164036),
(0.4000579116,	0.1916547891),
(0.4122855598,	0.2657302751),
(0.3897682267,	0.3744628158),
(0.3218496119,	0.1559752782),
(0.2959833070,	0.2812516245),
(0.4132174699,	0.2698100668),
(0.4010600790,	0.1988153634),
(0.3983719278,	0.1787862095),
(0.4134354098,	0.2854601690),
(0.4920112032,	0.1412368522),
(0.3958370784,	0.1564483425),
(0.2816607693,	0.2752504513),
(0.4155517978,	0.2960847052),
(0.3124918736,	0.1567548343),
(0.3936357530,	0.1305277271),
(0.3977453372,	0.1736695550),
(0.4065107464,	0.2337325768),
(0.3934827897,	0.1282315783),
(0.4514746238,	0.0973655800),
(0.3994542144,	0.1871787471),
(0.3466891157,	0.0826041253),
(0.3951370317,	0.1492265626),
(0.2704861637,	0.2617400496),
(0.2701991673,	0.2977922273),
(0.2719181145,	0.2227169845),
(0.3933545122,	0.1262431781),
(0.4239199957,	0.4890252375),
(0.4210931065,	0.4930983422),
(0.3202131331,	0.3028156498),
(0.0262032843,	0.6907471490)]
if __name__ == "__main__":
    run()
    time.sleep(10000)